#!/usr/bin/env python3
import os

from data_utils import *
from publish_utils import *
from kitti_util import *
from head_function import *
from tracklet_parser import parse_xml

DATA_PATH = '/home/jony/data/kitti/RawData/2011_09_26/2011_09_26_drive_0061_sync/'

if __name__ == '__main__':
    frame = 0
    rospy.init_node('kitti_node', anonymous=True)
    cam_pub = rospy.Publisher('kitti_cam', Image, queue_size=10)
    pcl_pub = rospy.Publisher('kitti_pcl', PointCloud2, queue_size=10)
    ego_pub = rospy.Publisher('kitti_ego_car', MarkerArray, queue_size=10)
    imu_pub = rospy.Publisher('kitti_imu', Imu, queue_size=10)
    gps_pub = rospy.Publisher('kitti_gps', NavSatFix, queue_size=10)
    box3d_pub = rospy.Publisher('kitti_3d', MarkerArray, queue_size=10)
    bridge = CvBridge()

    rate = rospy.Rate(10)

    tracklets = parse_xml('/home/jony/data/kitti/tracklets_0061/2011_09_26/2011_09_26_drive_0061_sync/tracklet_labels.xml')
    calib = Calibration('/home/jony/data/kitti/RawData/2011_09_26/', from_video=True)
    df_tracklets = get_frame_fields(tracklets)

    while not rospy.is_shutdown():
        df_tracklets_frame = df_tracklets[df_tracklets.frame==frame]

        types = np.array(df_tracklets_frame['type'])
        boxes_3d = np.array(df_tracklets_frame[['h', 'w', 'l', 'tx', 'ty', 'tz', 'rz']])
        boxes_2d = []
        corners_3d_velos = []
        for box_3d in boxes_3d:
            h, w, l = box_3d[0], box_3d[1], box_3d[2]
            tx, ty, tz = box_3d[3], box_3d[4], box_3d[5]
            rz = wrap_to_pi(box_3d[6])
            yaw = rz_to_ry(rz)
            # 将世界坐标系下的 tx, ty, tz 转换为相机坐标系下的 pos_x, pos_y, pos_z
            pos_x, pos_y, pos_z = convert_world_to_camera_coords(tx, ty, tz, calib)
            # 使用转换后的坐标计算3D框
            corners_2d, corners_3d_cam2 = compute_3d_box_cam2(h, w, l, pos_x, pos_y, pos_z, yaw, calib)
            corners_3d_velo = calib.project_rect_to_velo(corners_3d_cam2.T)
            box_2d = get_2d_boxes(corners_2d)
            boxes_2d.append(box_2d)
            corners_3d_velos.append(corners_3d_velo)

        image = read_camera(os.path.join(DATA_PATH, 'image_02/data/%010d.png' % frame))
        point_cloud = read_point_cloud(os.path.join(DATA_PATH, 'velodyne_points/data/%010d.bin' % frame))
        imu_data = read_imu(os.path.join(DATA_PATH, 'oxts/data/%010d.txt' % frame))

        publish_camera(cam_pub, bridge, image, boxes_2d, types)
        publish_point_cloud(pcl_pub, point_cloud)
        publish_ego_car(ego_pub)
        publish_imu(imu_pub, imu_data)
        publish_gps(gps_pub, imu_data)
        publish_3dbox(box3d_pub, corners_3d_velos, types)
        rospy.loginfo("Published!")
        rate.sleep()
        frame += 1
        frame %= 703
